"""
Case Type   : DATAVEC模块
Case Name   : 创建ivfpq 内积距离索引，索引选项pq_ksub参数覆盖，参数不同，执行时长比较
Create At   : 2025/4/14
Owner       : @zhanghuan
Description :
    step1:设置pq环境变量
    step2:修改参数enable_pq为true
    step3:创建3张表
    step4:创建自定义插入函数
    step5:表插入相同数据
    step6:创建索引，指定pq_ksub为1
    step7:创建索引，指定pq_ksub为128
    step8:创建索引，不指定pq_ksub，默认值256
    step9:不同pq_ksub值，创建索引时长比较
    step10:不同pq_ksub值，update数据
    step11:不同pq_ksub值索引，update操作时长比较
    step12:不同pq_ksub值，delete数据
    step13:不同pq_ksub值，insert数据
    step14:不同pq_ksub值索引，插入数据操作时长比较
Expect      :
    step1:设置pq环境变量;expect:成功
    step2:修改参数enable_pq为true;expect:成功
    step3:创建3张表;expect:成功
    step4:创建自定义插入函数;expect:成功
    step5:表插入相同数据;expect:成功
    step6:创建索引，指定pq_ksub为1;expect:成功
    step7:创建索引，指定pq_ksub为128;expect:成功
    step8:创建索引，不指定pq_ksub，默认值256;expect:成功
    step9:不同pq_ksub值，创建索引时长比较;expect:pq_ksub越大，执行时长越长
    step10:不同pq_ksub值，update数据;expect:pq_ksub越大，执行时长越长
    step11:不同pq_ksub值索引，update操作时长比较;expect:pq_ksub越大，执行时长越长
    step12:不同pq_ksub值，delete数据;expect:成功
    step13:不同pq_ksub值，insert数据;expect:成功
    step14:不同pq_ksub值索引，插入数据操作时长比较;expect:pq_ksub值越大，执行时长越长
History     :
"""
import os
import unittest

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node
from yat.test import macro

pri_dbuser = Node(node='PrimaryDbUser')
pq_lib_path = '/usr/local/sra_recall/lib/'
log = Logger()
log.info("----判断本地环境是否存在pq插件----")
check_cmd = f'if [ -d {pq_lib_path} ]; ' \
            f'then echo "True"; else echo "False"; fi'
log.info(check_cmd)
check_result = pri_dbuser.sh(check_cmd).result()
log.info(check_result)


@unittest.skipIf(check_result == 'False',
                 '----环境不具备PQ插件，跳过执行用例----')
class DataVecTestCase(unittest.TestCase):
    def setUp(self):
        log.info(f'----{os.path.basename(__file__)}:start----')
        self.sh = CommonSH('PrimaryDbUser')
        self.sta_1_sh = CommonSH('Standby1DbUser')
        self.sta_2_sh = CommonSH('Standby2DbUser')
        self.constant = Constant()
        self.com = Common()
        self.pq_env_name = 'DATAVEC_PQ_LIB_PATH'
        self.tb_name = 't_datavec_index_pq_case0096'
        self.id_name = 'id_datavec_index_pq_case0096'
        self.comsql = CommonSQL('PrimaryDbUser')
        self.fun_insert = 'f_datavec_index_pq_case0096'

    def test_main(self):
        step_txt = '----step3:创建3张表; expect:成功----'
        log.info(step_txt)
        create_tb = f'''set enable_indexscan=on;
        set enable_seqscan=off;
        drop table if exists {self.tb_name}_1,{self.tb_name}_2,\
        {self.tb_name}_3;
        drop index if exists {self.id_name}_1,{self.id_name}_2,{self.id_name}_3;
        create table {self.tb_name}_1 (id bigserial primary key, val \
        vector(320));
        create table {self.tb_name}_2 (id bigserial primary key, val \
        vector(320));
        create table {self.tb_name}_3 (id bigserial primary key, val \
        vector(320));'''
        result = self.sh.execut_db_sql(create_tb)
        log.info(result)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, result,
                      '执行失败:' + step_txt)

        step_txt = '----step4:创建自定义插入函数; expect:成功----'
        log.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        log.info(result)
        self.assertIn(self.fun_insert, result, '执行失败:' + step_txt)

        step_txt = '----step5:表插入相同数据; expect:成功----'
        log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(320,5000,'{self.tb_name}_1'," \
                     f"'val','vector');"
        result = self.sh.execut_db_sql(insert_sql)
        log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        result = self.sh.execut_db_sql(insert_sql)
        log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_2 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_3 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step6:创建索引，指定pq_ksub为1; expect:成功----'
        log.info(step_txt)
        create_index = (f"set maintenance_work_mem='100MB';"
                        f"create index {self.id_name}_1 on {self.tb_name}_1 "
                        f"using ivfflat (val vector_ip_ops) "
                        f"with(enable_pq=on,pq_ksub=1);")
        result = self.comsql.get_query_duration(create_index)
        duration1 = result[0]
        log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        step_txt = '----step7:创建索引，指定pq_ksub为128; expect:成功----'
        log.info(step_txt)
        create_index = (f"set maintenance_work_mem='100MB';"
                        f"create index {self.id_name}_2 on {self.tb_name}_2 "
                        f"using ivfflat (val vector_ip_ops) "
                        f"with(enable_pq=on,pq_ksub=128);")
        result = self.comsql.get_query_duration(create_index)
        log.info(result)
        duration2 = result[0]
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        step_txt = '----step8:创建索引，不指定pq_ksub，默认值256; expect:成功----'
        log.info(step_txt)
        create_index = (f"set maintenance_work_mem='100MB';"
                        f"create index {self.id_name}_3 on {self.tb_name}_3 "
                        f"using ivfflat (val vector_ip_ops) with(enable_pq=on);")
        result = self.comsql.get_query_duration(create_index)
        log.info(result)
        duration3 = result[0]
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        step_txt = ('----step9:不同pq_ksub值，创建索引时长比较; '
                    'expect:pq_ksub越大，执行时长越长----')
        log.info(step_txt)
        self.assertGreater(duration3, duration2, '执行失败:' + step_txt)
        self.assertGreater(duration2, duration1, '执行失败:' + step_txt)

        step_txt = '----step10:不同pq_ksub值，update数据; expect成功----'
        log.info(step_txt)
        sql = f"update {self.tb_name}_1 set val = val + " \
              f"array_fill(49,array[320])::vector where id>5000;"
        result = self.comsql.get_query_duration(sql)
        duration1 = result[0]
        log.info(result)
        self.assertIn(self.constant.UPDATE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)
        sql = f"update {self.tb_name}_2 set val = val + " \
              f"array_fill(49,array[320])::vector where id>5000;"
        result = self.comsql.get_query_duration(sql)
        duration2 = result[0]
        log.info(result)
        self.assertIn(self.constant.UPDATE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)
        sql = f"update {self.tb_name}_3 set val = val + " \
              f"array_fill(49,array[320])::vector where id>5000;"
        result = self.comsql.get_query_duration(sql)
        duration3 = result[0]
        log.info(result)
        self.assertIn(self.constant.UPDATE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        step_txt = ('----step11:不同pq_ksub值索引，update操作时长比较; '
                    'expect:pq_ksub越大，执行时长越长----')
        log.info(step_txt)
        self.assertGreater(duration3, duration2, '执行失败:' + step_txt)
        self.assertGreater(duration2, duration1, '执行失败:' + step_txt)

        step_txt = '----step12:不同pq_ksub值，delete数据; expect:功能正常----'
        log.info(step_txt)
        sql = f"delete from {self.tb_name}_1;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.DELETE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)
        sql = f"delete from {self.tb_name}_2;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.DELETE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)
        sql = f"delete from {self.tb_name}_3;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.DELETE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        step_txt = '----step13:不同pq_ksub值，insert数据; expect:成功----'
        log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(320,2000,'{self.tb_name}_1'," \
                     f"'val','vector');"
        result = self.comsql.get_query_duration(insert_sql)
        duration1 = result[0]
        log.info(result)
        self.assertIn(self.fun_insert, result[1], '执行失败:' + step_txt)
        insert_sql = f"select {self.fun_insert}(320,2000,'{self.tb_name}_2'," \
                     f"'val','vector');"
        result = self.comsql.get_query_duration(insert_sql)
        duration2 = result[0]
        log.info(result)
        self.assertIn(self.fun_insert, result[1], '执行失败:' + step_txt)
        insert_sql = f"select {self.fun_insert}(320,2000,'{self.tb_name}_3'," \
                     f"'val','vector');"
        result = self.comsql.get_query_duration(insert_sql)
        duration3 = result[0]
        log.info(result)
        self.assertIn(self.fun_insert, result[1], '执行失败:' + step_txt)

        step_txt = ('----step14:不同pq_ksub值索引，插入数据操作时长比较; '
                    'expect:pq_ksub值越大，执行时长越长----')
        log.info(step_txt)
        self.assertGreater(duration3, duration2, '执行失败:' + step_txt)
        self.assertGreater(duration2, duration1, '执行失败:' + step_txt)

    def tearDown(self):
        step_txt = '----this is teardown----'
        log.info(step_txt)
        step1_txt = '----teardown:清理测试数据; expect:成功----'
        log.info(step1_txt)
        drop_cmd = f"drop table if exists {self.tb_name}_1,{self.tb_name}_2," \
                   f"{self.tb_name}_3;"
        drop_result1 = self.sh.execut_db_sql(drop_cmd)
        log.info(drop_result1)
        drop_cmd = f"drop function if exists {self.fun_insert};"
        drop_result2 = self.sh.execut_db_sql(drop_cmd)
        log.info(drop_result2)

        log.info(f'----{os.path.basename(__file__)}:end----')
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, drop_result1,
                      '执行失败:' + step_txt + step1_txt)
        self.assertIn(self.constant.DROP_FUNCTION_SUCCESS_MSG, drop_result2,
                      '执行失败:' + step_txt + step1_txt)
